# sample_size_calibration_test.py
"""
专门用于样本量校准的测试脚本 - 修复类型错误版本
"""

import csv
import statistics
from datetime import datetime, timedelta
from collections import defaultdict
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor, as_completed
import multiprocessing as mp

import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

from osc05.optimization.get_cost import get_cost_with_simulation
from osc05.static.constants import Commons, Constants, DirectoryConfig

output_dir = Path(DirectoryConfig.OUTPUT_DIR)
output_dir.mkdir(parents=True, exist_ok=True)  # 自动创建目录（如果不存在）

def run_single_calibration(args):
    """运行单个校准测试 - 用于并行计算"""
    sample_size, calibration_seed, baseline_solution = args

    cost = get_cost_with_simulation(
        start_correction=baseline_solution,
        seed=calibration_seed,
        sample_size=sample_size,
        calibration_seed=calibration_seed,
        debug=False
    )

    return {
        'sample_size': sample_size,
        'calibration_seed': calibration_seed,
        'cost': cost
    }


class SampleSizeCalibrator:
    """样本量校准器 - 修复类型错误版本"""

    def __init__(self, threshold=0.1, max_workers=None):
        self.threshold = threshold
        self.max_workers = max_workers or min(mp.cpu_count(), 8)
        self.results = []
        self.cv_results = []

    def safe_float(self, value):
        """安全转换为浮点数"""
        try:
            return float(value)
        except (ValueError, TypeError):
            print(f"⚠️  警告: 无法转换值为浮点数: {value}, 类型: {type(value)}")
            return 0.0

    def run_calibration_test_parallel(self):
        """并行运行样本量校准测试"""
        baseline_solution = [0] * Constants.N_ACTIVITY
        self.results = []

        print("=" * 60)
        print("样本量校准测试 - 并行版本")
        print(f"使用 {self.max_workers} 个并行进程")
        print("=" * 60)

        # 准备所有测试任务
        tasks = []
        for sample_size in Commons.sample_sizes:
            for calibration_seed in Commons.seeds:
                tasks.append((sample_size, calibration_seed, baseline_solution))

        total_tasks = len(tasks)
        completed_tasks = 0

        print(f"开始执行 {total_tasks} 个校准任务...")
        start_time = datetime.now()

        # 使用进程池并行执行
        with ProcessPoolExecutor(max_workers=self.max_workers) as executor:
            # 提交所有任务
            future_to_task = {
                executor.submit(run_single_calibration, task): task
                for task in tasks
            }

            # 收集结果
            for future in as_completed(future_to_task):
                task = future_to_task[future]
                sample_size, calibration_seed, _ = task

                try:
                    result = future.result()
                    # 确保cost是数值类型
                    result['cost'] = self.safe_float(result['cost'])
                    self.results.append(result)
                    completed_tasks += 1

                    # 进度显示
                    progress = completed_tasks / total_tasks * 100
                    elapsed_time = (datetime.now() - start_time).total_seconds()
                    estimated_total = elapsed_time / (completed_tasks / total_tasks) if completed_tasks > 0 else 0
                    remaining = estimated_total - elapsed_time

                    print(f"[{completed_tasks}/{total_tasks}] {progress:.1f}% - "
                          f"样本量={sample_size}, 种子={calibration_seed}, "
                          f"成本={result['cost']:.2f}, "
                          f"预计剩余: {remaining / 60:.1f}分钟")

                except Exception as e:
                    print(f"任务失败: 样本量={sample_size}, 种子={calibration_seed}, 错误: {e}")

        total_time = (datetime.now() - start_time).total_seconds()
        print(f"\n所有任务完成! 总耗时: {total_time / 60:.1f} 分钟")
        print(f"平均每个任务: {total_time / total_tasks:.1f} 秒")

        # 后续分析 - 添加详细的错误处理
        try:
            print("开始数据分析...")
            self.analyze_results()
            self.save_calibration_results()
            self.generate_visualizations()
            optimal_size = self.determine_optimal_sample_size()
            return optimal_size
        except Exception as e:
            print(f"❌ 分析阶段出错: {e}")
            print("调试信息:")
            print(f"结果数量: {len(self.results)}")
            if self.results:
                print(f"第一个结果的cost类型: {type(self.results[0]['cost'])}")
                print(f"第一个结果的cost值: {self.results[0]['cost']}")
            if self.cv_results:
                print(f"CV结果数量: {len(self.cv_results)}")
                print(f"第一个CV结果的cv类型: {type(self.cv_results[0]['cv'])}")
            raise

    def analyze_results(self):
        """分析校准结果 - 彻底修复类型错误"""
        print("步骤1: 数据分组...")
        # 按样本量分组
        grouped = defaultdict(list)
        for result in self.results:
            # 确保cost是数值类型
            cost_value = self.safe_float(result['cost'])
            grouped[result['sample_size']].append(cost_value)

        print("步骤2: 计算统计量...")
        self.cv_results = []
        for sample_size in sorted(grouped.keys()):
            costs = grouped[sample_size]

            # 确保所有值都是数值类型
            costs = [self.safe_float(cost) for cost in costs]

            # 计算统计量
            try:
                mean_cost = statistics.mean(costs)
                std_cost = statistics.stdev(costs) if len(costs) > 1 else 0.0

                # 安全计算CV
                if mean_cost != 0:
                    cv = std_cost / abs(mean_cost)
                else:
                    cv = float('inf')

                self.cv_results.append({
                    'sample_size': sample_size,
                    'mean_cost': mean_cost,
                    'std_cost': std_cost,
                    'cv': cv,
                    'min_cost': min(costs),
                    'max_cost': max(costs),
                    'range': max(costs) - min(costs),
                    'n_repeats': len(costs)
                })

                print(f"  样本量 {sample_size}: 平均={mean_cost:.2f}, 标准差={std_cost:.2f}, CV={cv:.6f}")

            except Exception as e:
                print(f"❌ 计算样本量 {sample_size} 的统计量时出错: {e}")
                print(f"  成本数据: {costs}")

    def determine_optimal_sample_size(self):
        # 不要 try...except！
        self.cv_results.sort(key=lambda x: int(x['sample_size']))

        print("\n边际收益分析")
        optimal_size = self.cv_results[0]['sample_size']

        for i in range(1, len(self.cv_results)):
            prev = self.cv_results[i - 1]
            current = self.cv_results[i]

            # 直接打印类型（关键！）
            print(f"🔍 检查第{i}步:")
            print(f"  prev['cv'] = {repr(prev['cv'])} | type = {type(prev['cv'])}")
            print(f"  current['cv'] = {repr(current['cv'])} | type = {type(current['cv'])}")

            # 这一行就是报错行！让它崩溃，我们看 traceback
            improvement = (prev['cv'] - current['cv']) / prev['cv']

            if improvement < self.threshold:
                optimal_size = prev['sample_size']
                break
        else:
            optimal_size = self.cv_results[-1]['sample_size']

        return optimal_size

    def save_calibration_results(self):
        """保存校准测试结果"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

        # 保存详细结果
        detail_file = output_dir / f"sample_size_calibration_detail_{timestamp}.csv"
        with open(detail_file, 'w', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerow(['样本量', '校准种子', '基准方案成本'])
            for result in self.results:
                cost_value = self.safe_float(result['cost'])
                writer.writerow([result['sample_size'], result['calibration_seed'], f"{cost_value:.2f}"])

        # 保存CV分析结果
        cv_file = output_dir / f"sample_size_cv_analysis_{timestamp}.csv"
        with open(cv_file, 'w', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerow(['样本量', '平均成本', '标准差', '变异系数', '最小成本', '最大成本', '范围', '重复次数'])
            for result in self.cv_results:
                writer.writerow([
                    result['sample_size'],
                    f"{result['mean_cost']:.2f}",
                    f"{result['std_cost']:.2f}",
                    f"{result['cv']:.6f}",
                    f"{result['min_cost']:.2f}",
                    f"{result['max_cost']:.2f}",
                    f"{result['range']:.2f}",
                    result['n_repeats']
                ])

        print(f"\n校准结果已保存:")
        print(f"  详细结果: {detail_file}")
        print(f"  CV分析: {cv_file}")

    def save_marginal_analysis(self, improvement_data, optimal_size):
        """保存边际收益分析结果"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        marginal_file = output_dir / f"marginal_analysis_{timestamp}.csv"

        with open(marginal_file, 'w', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerow(['从样本量', '到样本量', '提升率', '前CV', '后CV', '是否选择'])

            for data in improvement_data:
                is_selected = (data['to_sample_size'] == optimal_size) or (
                        data['improvement'] < self.threshold and data['from_sample_size'] == optimal_size
                )
                writer.writerow([
                    data['from_sample_size'],
                    data['to_sample_size'],
                    f"{data['improvement']:.4f}",
                    f"{data['prev_cv']:.6f}",
                    f"{data['current_cv']:.6f}",
                    "是" if is_selected else "否"
                ])

            writer.writerow([])
            writer.writerow(['最优样本量', optimal_size])
            writer.writerow(['边际收益阈值', f"{self.threshold:.0%}"])

        print(f"  边际分析: {marginal_file}")

    def generate_visualizations(self):
        """生成可视化图表"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

        # 设置中文字体
        plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial']
        plt.rcParams['axes.unicode_minus'] = False

        # 创建图形
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 12))

        # 1. CV随样本量变化曲线
        sample_sizes = [r['sample_size'] for r in self.cv_results]
        cvs = [self.safe_float(r['cv']) for r in self.cv_results]

        ax1.plot(sample_sizes, cvs, 'bo-', linewidth=2, markersize=8)
        ax1.set_xscale('log')
        ax1.set_xlabel('样本量')
        ax1.set_ylabel('变异系数 (CV)')
        ax1.set_title('样本量与变异系数关系')
        ax1.grid(True, alpha=0.3)

        # 添加数据标签
        for i, (x, y) in enumerate(zip(sample_sizes, cvs)):
            ax1.annotate(f'{y:.4f}', (x, y), textcoords="offset points",
                         xytext=(0, 10), ha='center', fontsize=8)

        # 2. 成本分布箱线图
        cost_data = []
        labels = []
        for result in self.cv_results:
            sample_size = result['sample_size']
            costs = [self.safe_float(r['cost']) for r in self.results if r['sample_size'] == sample_size]
            cost_data.append(costs)
            labels.append(str(sample_size))

        ax2.boxplot(cost_data, labels=labels)
        ax2.set_xlabel('样本量')
        ax2.set_ylabel('基准方案成本')
        ax2.set_title('不同样本量下的成本分布')
        ax2.grid(True, alpha=0.3)
        plt.setp(ax2.xaxis.get_majorticklabels(), rotation=45)

        # 3. 边际收益分析
        if len(self.cv_results) > 1:
            improvements = []
            size_pairs = []
            for i in range(1, len(self.cv_results)):
                prev = self.cv_results[i - 1]
                current = self.cv_results[i]
                prev_cv = self.safe_float(prev['cv'])
                current_cv = self.safe_float(current['cv'])
                improvement = (prev_cv - current_cv) / prev_cv if prev_cv > 0 else 0
                improvements.append(improvement)
                size_pairs.append(f"{prev['sample_size']}\n→{current['sample_size']}")

            bars = ax3.bar(size_pairs, improvements, color='lightblue', alpha=0.7)
            ax3.axhline(y=self.threshold, color='red', linestyle='--',
                        label=f'阈值 ({self.threshold:.0%})')

            # 为超过阈值的柱子着色
            for i, (bar, imp) in enumerate(zip(bars, improvements)):
                if imp < self.threshold:
                    bar.set_color('lightcoral')

            ax3.set_xlabel('样本量变化')
            ax3.set_ylabel('精度提升率')
            ax3.set_title('边际收益分析')
            ax3.legend()
            ax3.grid(True, alpha=0.3)
            plt.setp(ax3.xaxis.get_majorticklabels(), rotation=45)

        # 4. 标准差随样本量变化
        stds = [r['std_cost'] for r in self.cv_results]
        ax4.plot(sample_sizes, stds, 'go-', linewidth=2, markersize=8)
        ax4.set_xscale('log')
        ax4.set_xlabel('样本量')
        ax4.set_ylabel('标准差')
        ax4.set_title('样本量与标准差关系')
        ax4.grid(True, alpha=0.3)

        plt.tight_layout()

        # 保存图形
        plot_file = output_dir / f"sample_size_analysis_{timestamp}.png"
        plt.savefig(plot_file, dpi=300, bbox_inches='tight')
        plt.close()

        print(f"  分析图表: {plot_file}")

    def run_calibration_test(self, use_parallel=True):
        """主运行函数"""
        if use_parallel:
            return self.run_calibration_test_parallel()
        else:
            return self.run_calibration_test_sequential()

    def run_calibration_test_sequential(self):
        """顺序运行版本（备用）"""
        baseline_solution = [0] * Constants.N_ACTIVITY
        self.results = []

        print("=" * 60)
        print("样本量校准测试 - 顺序版本")
        print("=" * 60)

        total_combinations = len(Commons.sample_sizes) * len(Commons.seeds)
        current_combination = 0
        start_time = datetime.now()

        for sample_size in Commons.sample_sizes:
            for calibration_seed in Commons.seeds:
                current_combination += 1

                cost = get_cost_with_simulation(
                    start_correction=baseline_solution,
                    seed=calibration_seed,
                    sample_size=sample_size,
                    calibration_seed=calibration_seed,
                    debug=False
                )

                # 确保cost是数值类型
                cost_value = self.safe_float(cost)

                self.results.append({
                    'sample_size': sample_size,
                    'calibration_seed': calibration_seed,
                    'cost': cost_value
                })

                # 进度显示
                progress = current_combination / total_combinations * 100
                elapsed_time = (datetime.now() - start_time).total_seconds()
                estimated_total = elapsed_time / (current_combination / total_combinations)
                remaining = estimated_total - elapsed_time

                print(f"[{current_combination}/{total_combinations}] {progress:.1f}% - "
                      f"样本量={sample_size}, 种子={calibration_seed}, "
                      f"成本={cost_value:.2f}, 预计剩余: {remaining / 60:.1f}分钟")

        total_time = (datetime.now() - start_time).total_seconds()
        print(f"\n所有任务完成! 总耗时: {total_time / 60:.1f} 分钟")

        # 后续分析
        self.analyze_results()
        self.save_calibration_results()
        self.generate_visualizations()
        optimal_size = self.determine_optimal_sample_size()

        return optimal_size


def main():
    """主函数"""
    calibrator = SampleSizeCalibrator(
        threshold=Commons.sample_size_threshold,
        max_workers=min(mp.cpu_count(), 8)
    )

    try:
        print(f"系统CPU核心数: {mp.cpu_count()}")
        print(f"使用的并行进程数: {calibrator.max_workers}")

        optimal_size = calibrator.run_calibration_test(use_parallel=True)

        print("\n" + "=" * 60)
        print("🎯 并行校准完成!")
        print("=" * 60)
        print(f"推荐的最优样本量: {optimal_size}")
        print(f"请在后线的GA优化中使用 sample_size={optimal_size}")
        print("=" * 60)

    except Exception as e:
        print(f"并行校准失败: {e}")
        print("尝试顺序执行...")

        try:
            optimal_size = calibrator.run_calibration_test(use_parallel=False)
            print(f"顺序执行完成，推荐样本量: {optimal_size}")
        except Exception as e2:
            print(f"顺序执行也失败: {e2}")
            print("请检查数据格式或文件权限")
            raise


if __name__ == '__main__':
    mp.freeze_support()
    main()